{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 生成自旋算符的函数\n",
    "def spin_operator_one_half():\n",
    "    op = dict()\n",
    "    op['i'] = np.eye(2)  # Identity\n",
    "    op['x'] = np.zeros((2, 2))\n",
    "    op['x'][0, 1] = 1 / 2\n",
    "    op['x'][1, 0] = 1 / 2\n",
    "    op['y'] = np.zeros((2, 2), dtype=np.complex)\n",
    "    op['y'][0, 1] = 1j / 2\n",
    "    op['y'][1, 0] = -1j / 2\n",
    "    op['z'] = np.zeros((2, 2))\n",
    "    op['z'][0, 0] = 1 / 2\n",
    "    op['z'][1, 1] = -1 / 2\n",
    "    return op\n",
    "\n",
    "\n",
    "# 生成二自旋海森堡哈密顿量\n",
    "def heisenberg_hamilt(j, h):\n",
    "    \"\"\"\n",
    "    :param j: list，耦合参数[Jx, Jy, Jz]\n",
    "    :param h: list，外磁场[hx, hy, hz]\n",
    "    :return H: 哈密顿量\n",
    "    \"\"\"\n",
    "    op = spin_operator_one_half()\n",
    "    H = j[0]*np.kron(op['x'], op['x']) + j[1]*np.kron(op['y'], op['y']) + \\\n",
    "        j[2]*np.kron(op['z'], op['z'])\n",
    "    H += h[0] * (np.kron(op['x'], op['i']) + np.kron(op['i'], op['x']))\n",
    "    H += h[1] * (np.kron(op['y'], op['i']) + np.kron(op['i'], op['y']))\n",
    "    H += h[2] * (np.kron(op['z'], op['i']) + np.kron(op['i'], op['z']))\n",
    "    if np.linalg.norm(np.imag(H)) < 1e-20:\n",
    "        H = np.real(H)\n",
    "    return H"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Spin-x operator = \n",
      "[[0.  0.5]\n",
      " [0.5 0. ]]\n",
      "\n",
      "Spin-y operator = \n",
      "[[ 0.+0.j   0.+0.5j]\n",
      " [-0.-0.5j  0.+0.j ]]\n",
      "\n",
      "Spin-z operator = \n",
      "[[ 0.5  0. ]\n",
      " [ 0.  -0.5]]\n"
     ]
    }
   ],
   "source": [
    "spin = spin_operator_one_half()\n",
    "print('Spin-x operator = ')\n",
    "print(spin['x'])\n",
    "print('\\nSpin-y operator = ')\n",
    "print(spin['y'])\n",
    "print('\\nSpin-z operator = ')\n",
    "print(spin['z'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hamilt = heisenberg_hamilt([1, 1, 1], [0, 0, 0])  # 无磁场的二自旋海森堡哈密顿量\n",
    "print('零磁场下海森堡哈密顿量 = ')\n",
    "print(hamilt)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
